from basic_component_related_management.models import *
import pandas as pd
from django.db import IntegrityError, transaction
import json
from tqdm import tqdm
import glob

def init_component_attribute():
    super_admin_entity = User.objects.filter(user_account="superadmin").first()
    # 初始化一些属性，全部都是short_string
    attribute_content = [
        {"name": "PCdata组分类别标识", "name_en": "nPCdataCompType"},
        {"name": "IUPAC名称", "name_en": "sIUPACName"},
        {"name": "PubChem标识", "name_en": "nPubChemID"},
        {"name": "DSI标识", "name_en": "nDSIID"},
        {"name": "DDB标识", "name_en": "nDDBID"},
        {"name": "DDBEAspen标识", "name_en": "sDDBEAspen"},
        {"name": "Aspen标识", "name_en": "sAspenID"},
        {"name": "TRC标识", "name_en": "nTRCID"},
        {"name": "Dippr标识", "name_en": "nDipprID"},
        {"name": "Open标识", "name_en": "nOpenID"},
    ]
    attribute_entity_list = []
    for item in attribute_content:
        attribute_entity = Attribute(
            name=item["name"],
            name_en = item["name_en"],
            data_type = "short_string",
            common_used=True,
            create_person = super_admin_entity,
            modify_person = super_admin_entity,
        )
        attribute_entity_list.append(attribute_entity)
    try:
        with transaction.atomic():
            Attribute.objects.bulk_create(attribute_entity_list)
    except IntegrityError as e:
        print(str(e))
    
def init_component_data():

    attribute_content = [
        {"name": "PCdata组分类别标识", "name_en": "nPCdataCompType"},
        {"name": "IUPAC名称", "name_en": "sIUPACName"},
        {"name": "PubChem标识", "name_en": "nPubChemID"},
        {"name": "DSI标识", "name_en": "nDSIID"},
        {"name": "DDB标识", "name_en": "nDDBID"},
        {"name": "DDBEAspen标识", "name_en": "sDDBEAspen"},
        {"name": "Aspen标识", "name_en": "sAspenID"},
        {"name": "TRC标识", "name_en": "nTRCID"},
        {"name": "Dippr标识", "name_en": "nDipprID"},
        {"name": "Open标识", "name_en": "nOpenID"},
    ]

    super_admin_entity = User.objects.filter(user_account="superadmin").first()
    # 获取所有的属性内容并生成name_en对entity的map
    attribute_name_en_entity_map = {}
    attribute_entity_list = Attribute.objects.all()
    for attribute_entity in attribute_entity_list:
        attribute_name_en_entity_map[attribute_entity.name_en] = attribute_entity
    matched_files = glob.glob("../final_basic_tables/A074. PCdata基表-2-组分表*.xlsx")
    latest_file = max(matched_files, key=lambda x: x.split("-")[-1]) if matched_files else None
    sheet_name = "Sheet1"
    df = pd.read_excel(latest_file,sheet_name=sheet_name,header=0,skiprows=[0], dtype=str)
    start_row = 0
    num_columns = 20
    selected_data = df.iloc[start_row:,:num_columns].fillna("")
    # print(selected_data[0:10])
    block_size = 1000


    blocks = []
    current_block = []
    alias_clocks = []
    current_alias_clock = []
    attribute_blocks = []
    current_attribute_block = []

    # print(len(selected_data))
    # print(selected_data)
    # return 

    for index,row in tqdm(selected_data.iterrows(), desc="Processing component data"):
    # for index,row  in selected_data.iterrows():
        component_name_en = ""
        sCompName = row["sCompName"]
        alias = row["sCompNameAliasList"].split(";")
        sCompNameCN = row["sCompNameCN"]
        sIUPACName = row["sIUPACName"]
        sFormula = row["sFormula"]
        sCASNo = row["sCASNo"]
        sSMILES = row["sSMILES"]
        # 获取一个英文名称
        for name in [sCompNameCN, sCompName, sIUPACName, sFormula, sCASNo, sSMILES]:
            if name != "":
                component_name_en = name
                break
        # 构建一个组分
        component_entity = ComponentIdentifier(
            name=sCompNameCN,
            name_en = component_name_en,
            molecular_formula = sFormula,
            sei_id = row["nPCdataCompID"],
            cas_number = sCASNo,
            smiles_number = sSMILES,
            inchi_number = row["sInChiKey"],
            create_person = super_admin_entity,
            modify_person = super_admin_entity
        )
        current_block.append(component_entity)
        # 构建别名
        for alia in alias:
            if alias!="":
                current_alias_clock.append(
                    ComponentIdentifierAlias(
                        componentidentifier = component_entity,
                        alias = alia
                    )
                )
        # 构建属性值
        for attribute_item in attribute_content:
            if row[attribute_item["name_en"]]!="":
                current_attribute_block.append(
                    ComponentAttributeValue(
                        componentidentifier = component_entity,
                        attribute = attribute_name_en_entity_map[attribute_item["name_en"]],
                        attribute_value_char = row[attribute_item["name_en"]]
                    )
                )
        # 如果大于等于了block_size，则分块
        if len(current_block) >=block_size:
            blocks.append(current_block)
            alias_clocks.append(current_alias_clock),
            attribute_blocks.append(current_attribute_block)
            current_block = []
            current_alias_clock = []
            current_attribute_block = []

    if current_block!=[]:
        blocks.append(current_block)
        alias_clocks.append(current_alias_clock),
        attribute_blocks.append(current_attribute_block)

    block_num = len(blocks)
    for index in range(block_num):
        if index <-1:
            continue
        try:
            with transaction.atomic():
                ComponentIdentifier.objects.bulk_create(
                    blocks[index]
                )
                ComponentIdentifierAlias.objects.bulk_create(
                    alias_clocks[index]
                )
                ComponentAttributeValue.objects.bulk_create(
                    attribute_blocks[index]
                )
                print("录入块:", index)
                # raise Exception("test")
        except IntegrityError as e:
            print(str(e))
            print("error in block:",index)
            return 
    return ComponentIdentifier.objects.count()


def init_component_content():
    print("开始录入组分")
    init_component_attribute()
    total_component = init_component_data()
    print(f"组分录入完成，共录入：{total_component}")
